import argparse
import matplotlib.pyplot as plt
import numpy as np
import csv
from scipy.stats import gaussian_kde

def parse_arguments():
    parser = argparse.ArgumentParser(description="Plot methylation level distribution.")
    parser.add_argument("--input", required=True, help="Path to the input TSV file.")
    parser.add_argument("--depth", type=int, default=10, help="Minimum sequencing depth for filtering (default: 10).")
    parser.add_argument("--output", required=True, help="Path to save the output plot.")
    return parser.parse_args()

def load_data(file_path, min_depth):
    methylation_levels = []
    with open(file_path, "r") as file:
        reader = csv.reader(file, delimiter="\t")
        for row in reader:
            try:
                depth = int(row[9])
                if depth >= min_depth:
                    methylation_level = float(row[10]) #* 100
                    if methylation_level > 0:  # Exclude methylation level of 0
                        methylation_levels.append(methylation_level)
            except (ValueError, IndexError):
                continue
    return methylation_levels

def plot_distribution(methylation_levels, output_path):
    kde = gaussian_kde(methylation_levels, bw_method='scott')
    x = np.linspace(51, 100, 1000)  # Start from 1 to exclude 0
    density = kde(x)

    plt.fill_between(x, density, alpha=0.5)
    plt.xlabel("Methylation Level (%)")
    plt.ylabel("Density")
    plt.title("Methylation Level Distribution")
    plt.xticks(np.arange(50, 101, 10))
    plt.grid(axis="y", linestyle="--", alpha=0.7)
    plt.savefig(output_path)
    plt.close()

def main():
    args = parse_arguments()
    methylation_levels = load_data(args.input, args.depth)
    if not methylation_levels:
        print("No data to plot. Please check the input file or filtering conditions.")
        return
    plot_distribution(methylation_levels, args.output)

if __name__ == "__main__":
    main()
